#! /usr/bin/python
"""
This class Itol takes and uploads variables to the iTOL server.  
The most basic way to run this program is to call it from command line
with the path to a tree file as the argument
The other way is to import and run in another Python program
to upload a tree is to use either the setTreeFile or 
setTreeContents methods to set the tree then call submitFile.  
If iTOL accepts the tree, submitFile will return True and
the treeID will be found by the getTreeID method.
See the README.txt for details
@author: Albert Wang (albertyw@mit.edu)
@date: March 7, 2010
"""
import os
import sys
import ItolFile
import Comm
import ItolExport

class Itol:
    """
    Create the Itol Class
    """
    def __init__(self):
        self.params = dict({})
        self.export_params = dict({})
        self.datasets = []
        self.comm = Comm.Comm()
    
    ###Setting the Tree
    def set_tree_location(self, tree_location):
        """
        set the tree file location (use setTreeContents to use tree data instead of a tree file)
        @param treeLocation: the path to the tree file
        """
        tree_file = ItolFile.ItolFile()
        tree_file.set_file_location(tree_location)
        self.params['treeFile'] = tree_file
    
    def set_tree_contents(self, tree_contents):
        """
        Set the tree contents
        @param treeContents: the structure of the tree to be uploaded
        """
        tree_file = ItolFile.ItolFile()
        tree_file.set_file_contents(tree_contents)
        self.params['treeFile'] = tree_file
    
    def get_tree_contents(self):
        """
        Get the tree contents
        @return: the structure of the tree
        """
        return self.params['treeFile'].get_file_contents()
    
    ###Setting and Getting Parameters
    def add_upload_param_dict(self, params):
        """
        Set a dictionary of params for iTOL.  treeFile should be a string of the file location
        @param params: a dictionary of the parms to be uploaded to iTOL
        """
        self.params.update(params)
    
    def set_upload_param_value(self, var_name, var_value):
        """
        Set a single param for iTOL.  A list of params can be found at http://itol.embl.de/help/batch_help.shtml
        @param varName: the name of the parameter (e.g. treeName)
        @param varValue: the value of the parameter (e.g. tree5)
        """
        self.params[var_name] = var_value
    
    def add_dataset(self, dataset):
        """
        Add a dataset to the tree
        @param dataset: an instance of the dataset class from dataset.py
        """
        self.datasets.append(dataset)
    
    def get_params(self):
        """
        Get the parameters to be uploaded with the tree
        @return: dictionary of parameters
        """
        return self.params
    
    ###Uploading Tree
    def upload_tree(self):
        """
        Prepare the parameters then
        Upload the tree and other parameters to the iTOL server
        @return: whether the upload is good
        """
        #Prepare params, open tree File
        upload_params = self.params.copy()
        upload_params['treeFile'] = upload_params['treeFile'].get_file_handle()
        #Attach datasets
        for i in range(len(self.datasets)):
            dataset = self.datasets[i]
            if not dataset.ready_upload():
                print 'Dataset parameters not ready'
                return False
            dataset = dataset.get_params_upload(i+1)
            upload_params.update(dataset)
        if upload_params['treeFile'] == '':
            print 'No File has been selected.  \
            Use treeFile() to set the tree file location'
            return False
        good_upload = self.comm.upload_tree(upload_params)
        return good_upload
        
    ###Reading the output
    def upload_output(self):
        """
        @return the itol output of the upload
        """
        return self.comm.upload_output
    
    def tree_id(self):
        """
        Returns the uploaded tree ID.  Go to http://itol.embl.de/upload.cgi and enter the tree ID into the "Retrieve a previously uploaded tree"
        @return the treeID of the uploaded tree
        """
        return self.comm.tree_id
    
    def get_webpage(self):
        """
        Returns a webpage that will redirect to the tree on itol (the itol website doesn't have direct urls to trees
        @return: website to view the tree online
        """
        webpage = "http://itol.embl.de/external.cgi?tree="+\
            str(self.tree_id())+"&restore_saved=1"
        return webpage
    
    def get_itol_export(self):
        """
        Returns an instance of ItolExport in preparation of exporting from the generated tree
        @return: instance of ItolExport
        """
        itol_exporter = ItolExport.ItolExport()
        itol_exporter.set_export_param_value('tree', self.tree_id())
        return itol_exporter

def print_help():
    """
    Prints usage info to command line
    """
    print 'Usage: python itol.py TREEFILE [-r]'
    print '  or:  python ito.py -h'
    print 'TREEFILE is the path to the tree file to be uploaded to iTOL'
    print 'Calling itol.py with an acceptable tree file \
    will return a URL to access the tree'
    print ''
    print '  -h    Display this help message'
    print '  -r    Verbose Output'
    print ''
    print 'Use itolexport.py to export uploaded trees from iTOL'
    print 'Read the README and itol.py source code for \
    information about more powerful iTOL calls.'
    print ''
    print 'Report bugs to <albertyw@mit.edu>'

if __name__ == "__main__":
    """
    This will run if itol.py is run by itself instead of imported
    @param sys.argv
    """
    #Check For Arguments
    if len(sys.argv) == 1:
        print 'Use -h for help'
        sys.exit(0)
    arg = sys.argv[1]
    #Help Argument
    if arg == '--help' or arg == '-h':
        print_help()
        sys.exit(0)
    #Check Verbosity
    verbose = False
    if len(sys.argv) == 3:
        verbose = True
    #Check if argument is a path
    if not os.path.isfile(arg):
        print arg+' is not a file.'
        print 'Use -h for help'
        sys.exit(0)
    tree_location = arg
    #Submit the iTOL file and output a web page to view the tree
    if verbose:
        print 'Submitting tree to server...'
    itol = Itol()
    itol.set_tree_location(tree_location)
    submit_good = itol.upload_tree()
    if submit_good:
        print 'Good Upload.  URL: '
        print itol.get_webpage()
    else:
        print 'Bad Upload Error: '
        print itol.upload_output
    sys.exit(0)
